package com.chatplus.application.processor.handler;

import com.chatplus.application.common.constant.RedisPrefix;
import com.chatplus.application.common.domain.response.SensitiveWordFilterResultResponse;
import com.chatplus.application.common.enumeration.UserStatusEnum;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.common.util.SpringCtxUtils;
import com.chatplus.application.domain.dto.ChatConfigDto;
import com.chatplus.application.domain.dto.UserCustomChatConfigDto;
import com.chatplus.application.domain.entity.account.UserEntity;
import com.chatplus.application.domain.entity.chat.ChatHistoryEntity;
import com.chatplus.application.domain.entity.chat.ChatModelEntity;
import com.chatplus.application.domain.entity.chat.ChatRoleEntity;
import com.chatplus.application.processor.channel.chat.IChatAiProcessorService;
import com.chatplus.application.processor.constant.AIConstants;
import com.chatplus.application.processor.constant.AiPlatformConstants;
import com.chatplus.application.processor.handler.dto.ChatRecordMessage;
import com.chatplus.application.processor.handler.dto.ReplyChatMessage;
import com.chatplus.application.service.account.UserService;
import com.chatplus.application.service.basedata.ConfigService;
import com.chatplus.application.service.basedata.SensitiveWordService;
import com.chatplus.application.service.chat.ChatHistoryService;
import com.chatplus.application.service.chat.ChatModelService;
import com.chatplus.application.service.chat.ChatRoleService;
import com.fasterxml.jackson.core.JsonProcessingException;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.redisson.api.RBucket;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import com.chatplus.application.processor.AIChannelServiceProvider;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * AI聊天Socket会话处理器
 * 基于Spring WebSocket
 */
@Component
public class ChatAiSocketHandler extends AbstractWebSocketHandler {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(ChatAiSocketHandler.class);
    //在线总数
    private static int onlineCount;
    //chatId -目前是按浏览器随机生成
    private String chatId;

    private String platform;

    private static CopyOnWriteArraySet<ChatAiSocketHandler> webSocketSet = new CopyOnWriteArraySet<>();
    /**
     * 使用map对象，便于根据userId来获取对应的WebSocket，或者放redis里面
     */
    private static final Map<String, ChatAiSocketHandler> webSocketMap = new ConcurrentHashMap<>();
    /**
     * 为了保存在线用户信息，在方法中新建一个list存储一下【实际项目依据复杂度，可以存储到数据库或者缓存】
     */
    private static final List<WebSocketSession> SESSIONS = Collections.synchronizedList(new ArrayList<>());
    private final ChatModelService chatModelService;
    private final ChatRoleService chatRoleService;
    private final ChatHistoryService chatHistoryService;
    private final UserService userService;
    private final SensitiveWordService sensitiveWordService;
    private final ConfigService configService;
    private final RedissonClient redissonClient;

    @Autowired
    public ChatAiSocketHandler(ChatModelService chatModelService,
                               ChatRoleService chatRoleService,
                               ChatHistoryService chatHistoryService,
                               UserService userService,
                               SensitiveWordService sensitiveWordService,
                               ConfigService configService,
                               RedissonClient redissonClient) {
        this.chatModelService = chatModelService;
        this.chatRoleService = chatRoleService;
        this.chatHistoryService = chatHistoryService;
        this.userService = userService;
        this.sensitiveWordService = sensitiveWordService;
        this.configService = configService;
        this.redissonClient = redissonClient;
    }

    /**
     * 获取当前连接数
     */
    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    /**
     * 当前连接数加一
     */
    public static synchronized void addOnlineCount() {
        ChatAiSocketHandler.onlineCount++;
    }

    /**
     * 当前连接数减一
     */
    public static synchronized void subOnlineCount() {
        ChatAiSocketHandler.onlineCount--;
    }

    /**
     * 连接成功后
     */
    @Override
    public void afterConnectionEstablished(@NotNull WebSocketSession session) throws JsonProcessingException {
        this.chatId = session.getAttributes().get("chat_id").toString();
        webSocketSet.add(this);
        SESSIONS.add(session);
        if (webSocketMap.containsKey(chatId)) {
            webSocketMap.remove(chatId);
            webSocketMap.put(chatId, this);
        } else {
            webSocketMap.put(chatId, this);
            addOnlineCount();
        }
        // ------------------------初始化校验------------------------
        // 1. 判断模型是否启用
        String modelId = session.getAttributes().get("model_id").toString();
        ChatModelEntity chatModelEntity = chatModelService.getById(Long.valueOf(modelId));
        if (chatModelEntity == null || !chatModelEntity.getEnabled()) {
            replyErrorMessage(session, "当前AI模型暂未启用，连接已关闭！！！");
            optClose(session);
            return;
        }
        platform = chatModelEntity.getPlatform();
        session.getAttributes().put("model_weight", chatModelEntity.getWeight());
        session.getAttributes().put("platform", platform);
        session.getAttributes().put("modelName", chatModelEntity.getValue());
        // 2. 判断角色是否启用
        String roleId = session.getAttributes().get("role_id").toString();
        ChatRoleEntity chatRoleEntity = chatRoleService.getById(Long.valueOf(roleId));
        if (chatRoleEntity == null || !chatRoleEntity.getEnable()) {
            replyErrorMessage(session, "当前聊天角色不存在或者未启用，连接已关闭！！！");
            optClose(session);
            return;
        }
        Long userId = Long.valueOf(session.getAttributes().get("userId").toString());
        // 3. 对话上下文加载
        ChatConfigDto chatConfigDto = configService.getChatConfig();
        if (chatConfigDto == null) {
            replyErrorMessage(session, "加载系统配置失败，连接已关闭！！！");
            optClose(session);
            return;
        }
        if (chatConfigDto.getContextDeep() > 0 && chatConfigDto.isEnableContext()) {
            List<ChatHistoryEntity> historyList = chatHistoryService.getChatHistoryByUserId(userId, null, chatId);
            Map<Long, String> promptMap = historyList.stream()
                    .limit(chatConfigDto.getContextDeep() * 2L) // 取前十条记录
                    .filter(history -> history.getType().equals("prompt")) // 过滤掉 reply 类型的记录
                    .sorted(Comparator.comparing(ChatHistoryEntity::getCreatedAt)) // 再次根据时间戳升序排序
                    .collect(ConcurrentHashMap::new,
                            (m, v) -> m.put(v.getItemId(), v.getContent()),
                            ConcurrentHashMap::putAll);
            List<ChatRecordMessage> recordMessageList =
                    historyList.stream()
                            .sorted(Comparator.comparing(ChatHistoryEntity::getId))
                            .filter(history -> history.getType().equals("reply"))
                            .map(history -> {
                                String prompt = promptMap.get(history.getItemId());
                                String reply = history.getContent();
                                if (StringUtils.isNotEmpty(prompt) && StringUtils.isNotEmpty(reply)) {
                                    ChatRecordMessage chatRecordMessage = new ChatRecordMessage();
                                    chatRecordMessage.setPrompt(prompt);
                                    chatRecordMessage.setReply(reply);
                                    return chatRecordMessage;
                                }
                                return null;
                            })
                            .filter(Objects::nonNull)
                            .collect(ArrayList::new, ArrayList::add, ArrayList::addAll);

            redissonClient.getBucket(RedisPrefix.TURN_RIGHT_REDIS_PREFIX + chatId)
                    .set(recordMessageList, Duration.ofMinutes(30));
        }

        // 会话上下文处理
        LOGGER.message("建立连接").context("连接ID", this.chatId).context("当前连接数", Optional.of(getOnlineCount())).info();
    }

    /**
     * 处理发送来的消息
     */
    @Override
    public void handleMessage(@NotNull WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        LOGGER.message("收到消息").context("连接ID", this.chatId)
                .context("参数列表", session.getAttributes())
                .context("消息", message.getPayload()).info();
        String payload = (String) message.getPayload();
        // ------------------------用户信息校验------------------------
        String verifyResult = verifyUserRequest(session, payload);
        if (StringUtils.isNotEmpty(verifyResult)) {
            replyErrorMessage(session, verifyResult);
            replyErrorMessage(session, AIConstants.ERR_IMG);
            return;
        }
        // ------------------------用户校验结束------------------------
        //会话上下文处理
        try {
            RBucket<List<ChatRecordMessage>> rBucket = redissonClient.getBucket(RedisPrefix.TURN_RIGHT_REDIS_PREFIX + chatId);
            List<ChatRecordMessage> recordMessageList = rBucket.isExists() ? rBucket.get() : new ArrayList<>();
            // 最大只能有十个上下文
            if (recordMessageList.size() >= 10) {
                recordMessageList = recordMessageList.subList(1, 10);
            }
            // 当前本次的答题
            ChatRecordMessage chatRecordMessage = new ChatRecordMessage();
            chatRecordMessage.setPrompt(payload);
            recordMessageList.add(chatRecordMessage);

            AIChannelServiceProvider aiChannelServiceProvider = SpringCtxUtils.getBean(AIChannelServiceProvider.class);
            IChatAiProcessorService aiProcessor = aiChannelServiceProvider.getPayChannelService(platform);
            aiProcessor.setSession(session);
            aiProcessor.process(recordMessageList);
        } catch (Exception e) {
            replyErrorMessage(session, "服务端发生异常，等待管理员修复！");
            replyErrorMessage(session, AIConstants.ERR_IMG);
            LOGGER.message("发送消息异常").exception(e).error();
        }
    }

    private String verifyUserRequest(WebSocketSession session, String payload) {
        // 敏感字处理
        SensitiveWordFilterResultResponse wordFilterResultResponse = sensitiveWordService.filter(payload);
        if (wordFilterResultResponse.isContainsSensitiveWord()) {
            LOGGER.message("用户输入包含敏感词").context("input", payload)
                    .context("敏感词", wordFilterResultResponse.getSensitiveWordList()).info();
            return "您的输入包含敏感词，请重新输入！如果触发误报，请联系管理员！";
        }
        // TODO 以下所有的逻辑有待验证
        Long userId = Long.valueOf(session.getAttributes().get("userId").toString());
        UserEntity userEntity = userService.getById(userId);
        if (userEntity == null) {
            return "非法用户，请联系管理员！";
        }
        if (userEntity.getStatus() != UserStatusEnum.OK) {
            return "您的账号已经被禁用，如果疑问，请联系管理员！";
        }
        Integer modelWeight = Integer.valueOf(session.getAttributes().get("model_weight").toString());
        Integer calls = userEntity.getCalls();
        if (calls < modelWeight) {
            return String.format("您当前剩余对话次数（%d）已不足以支付当前模型的单次对话需要消耗的对话额度（%d）！"
                    , calls, modelWeight);
        }
        UserCustomChatConfigDto chatConfig = userEntity.getChatConfig();
        UserCustomChatConfigDto.ApiKeysBean apiKeysBean = chatConfig.getApiKeys();
        String apiKey = null;
        switch (platform) {
            case AiPlatformConstants.OPEN_AI:
                apiKey = apiKeysBean.getOpenAI();
                break;
            case AiPlatformConstants.CHAT_GLM:
                apiKey = apiKeysBean.getChatGLM();
                break;
            case AiPlatformConstants.Azure:
                apiKey = apiKeysBean.getAzure();
                break;
            default:
                break;
        }
        if (calls <= 0 && StringUtils.isEmpty(apiKey)) {
            return "您的对话次数已经用尽，请联系管理员或者充值点卡继续对话！";
        }
        return null;
    }

    @Override
    public void handleTransportError(WebSocketSession session, @NotNull Throwable exception) {
        LOGGER.message("连接发送错误").context("连接信息", session.getAttributes()).context("exception", exception.getMessage()).error();
        replyErrorMessage(session, "连接发送错误,等待管理员修复！");
    }

    @Override
    public void afterConnectionClosed(@NotNull WebSocketSession session, @NotNull CloseStatus closeStatus) {
        webSocketSet.remove(this);
        if (webSocketMap.containsKey(chatId)) {
            webSocketMap.remove(chatId);
            subOnlineCount();
        }
        LOGGER.message("断开连接").context("连接ID", this.chatId).context("当前连接数", getOnlineCount()).info();
    }

    public static void sendMessage(WebSocketSession session, String message) {
        if (session == null || !session.isOpen()) {
            LOGGER.message("session会话已经关闭").context("message", message).error();
            return;
        }
        try {
            if (AIConstants.END_LABEL.equals(message)) {
                ReplyChatMessage end = new ReplyChatMessage("end", null);
                session.sendMessage(new BinaryMessage(end.toByte()));
                return;
            }
            if (AIConstants.START_LABEL.equals(message)) {
                ReplyChatMessage start = new ReplyChatMessage("start", null);
                session.sendMessage(new BinaryMessage(start.toByte()));
                return;
            }
            ReplyChatMessage chunk = new ReplyChatMessage("middle", message);
            session.sendMessage(new BinaryMessage(chunk.toByte()));
        } catch (Exception e) {
            LOGGER.message("发送消息异常").exception(e).error();
        }
    }

    /**
     * 操作进行下线
     */
    public void optClose(WebSocketSession session) {
        // 判断当前连接是否还在线
        if (session.isOpen()) {
            try {
                // 关闭连接
                session.close(CloseStatus.NORMAL);
            } catch (Exception e) {
                LOGGER.message("关闭session异常").exception(e).error();
            }

        }
    }

    public static void replyErrorMessage(WebSocketSession session, String message) {
        sendMessage(session, AIConstants.START_LABEL);
        sendMessage(session, message);
        sendMessage(session, AIConstants.END_LABEL);
    }
}
